'''
This script shows how to predict stock prices using a basic RNN
'''
import tensorflow as tf
import numpy as np
import matplotlib as plt
import os

tf.set_random_seed(777)  #重现数据 reproducibility

import matplotlib.pyplot as plt

def MinMaxScaler(data):
    '''
    # 1.Min-Max Normalization 将原始数据线性化的方法转换到[0 1]的范围 :
    # 2.将原始数据集归一化为均值为0、方差1的数据集:
    '''
    numerator = data - np.min(data, 0)
    denominator = np.max(data, 0) - np.min(data, 0)
    #  + 1e-7 是 在分母项加上一个 eplison项 主要作用 防止分母为零
    return numerator / (denominator + 1e-7)

seq_length = 7 #序列长度
data_dim = 5 #数据（特征值）
hidden_dim = 10  #隐藏单元个数
output_dim = 1 #输出
learning_rate = 0.01 #学习率
iterations = 100#500  迭代次数

xy = np.loadtxt(r'../../../../large_data/DL1/stock/data-02-stock_daily.csv', delimiter=',')
xy = xy[::-1]  #逆序后按正常时间顺序reverse order (chronically ordered)
xy = MinMaxScaler(xy)
x = xy
y = xy[:, [-1]]  # 把收盘价做为实际值标签 Close as label

# 构建数据集
dataX = []
dataY = []  #实际值（标签）
for i in range(0, len(y) - seq_length):
    _x = x[i:i + seq_length]
    _y = y[i + seq_length]  #下一个收盘价 Next close price
    # print(_x, "->", _y)
    dataX.append(_x)
    dataY.append(_y)

# 切分训练集和测试集 train/test split
train_size = int(len(dataY) * 0.7)
test_size = len(dataY) - train_size
trainX, testX = np.array(dataX[0:train_size]), np.array(dataX[train_size:len(dataX)])
trainY, testY = np.array(dataY[0:train_size]), np.array(dataY[train_size:len(dataY)])

X = tf.placeholder(tf.float32, [None, seq_length, data_dim])  #三维数据，样本数、序列长度、输入数据大小
Y = tf.placeholder(tf.float32, [None, 1]) #标签收盘价

# 建立LSTM网络build a LSTM network
# cell = tf.contrib.rnn.BasicLSTMCell(num_units=hidden_dim, state_is_tuple=True, activation=tf.tanh)
cell = tf.contrib.rnn.LSTMCell(num_units=hidden_dim, state_is_tuple=True, activation=tf.tanh)
# 输入数据inputs(batch_size, time_steps, input_size) 三维数据，样本数、序列长度、输入数据大小
# outputs是time_steps步里所有的输出(batch_size, time_steps, cell.output_size）
# state 是最后一步的隐状态（batch_size, cell.state_size）
outputs, _states = tf.nn.dynamic_rnn(cell, X, dtype=tf.float32)
# 全连接层。用最后cell的输出 We use the last cell's output
Y_pred = tf.contrib.layers.fully_connected(outputs[:, -1], output_dim, activation_fn=None)

# 代价/损失值cost/loss
loss = tf.reduce_sum(tf.square(Y_pred - Y))  #均方差用于训练集 sum of the squares
train = tf.train.AdamOptimizer(learning_rate).minimize(loss)
# 均方根误差(RMSE) 用于测试集算准确率
targets = tf.placeholder(tf.float32, [None, 1])
predict = tf.placeholder(tf.float32, [None, 1])
rmse = tf.sqrt(tf.reduce_mean(tf.square(targets - predict)))  #均方根
r2 = 1. - tf.reduce_mean(tf.square(targets - predict)) / tf.reduce_mean(tf.square(targets - tf.reduce_mean(targets)))
loss_s = tf.summary.scalar('loss', loss)
summary = tf.summary.merge([loss_s])
rmse_s = tf.summary.scalar('rmse', rmse)
r2_s = tf.summary.scalar('r2', r2)
summary_calc = tf.summary.merge([rmse_s, r2_s])

with tf.Session() as sess:#创建会话
    with tf.summary.FileWriter('./_log/' + os.path.basename(__file__), sess.graph) as fw:
        sess.run(tf.global_variables_initializer()) #全局变量初始化
        for i in range(iterations):
            _, step_loss, sv = sess.run([train, loss, summary], feed_dict={X: trainX, Y: trainY})
            fw.add_summary(sv, i)
            if i % 10 == 0:
                test_predict = sess.run(Y_pred, feed_dict={X: testX})
                rmse_val = sess.run(rmse, feed_dict={targets: testY, predict: test_predict})
                r2_val = sess.run(r2, feed_dict={targets: testY, predict: test_predict})
                print("[step: {}] loss: {} rmse: {} r2: {}".format(i, step_loss, rmse_val, r2_val))
                sv = sess.run(summary_calc, feed_dict={targets: testY, predict: test_predict})
                fw.add_summary(sv, i)

    # test
    test_predict = sess.run(Y_pred, feed_dict={X: testX})
    rmse_val = sess.run(rmse, feed_dict={targets: testY, predict: test_predict})
    r2_val = sess.run(r2, feed_dict={targets: testY, predict: test_predict})
    print("RMSE: {}, R2: {}".format(rmse_val, r2_val))

    # 画图
    plt.plot(testY) #蓝色实际值
    plt.plot(test_predict, 'r')  # 红色预测值
    plt.xlabel("Time Period")
    plt.ylabel("Stock Price")
    plt.show()